// Physical memory allocator, for user processes,
// kernel stacks, page-table pages,
// and pipe buffers. Allocates whole 4096-byte pages.

#include "types.h"
#include "param.h"
#include "memlayout.h"
#include "spinlock.h"
#include "riscv.h"
#include "defs.h"

void freerange(void *pa_start, void *pa_end);

extern char end[]; // first address after kernel.
                   // defined by kernel.ld.

struct run {
  struct run *next;
};

struct {
  struct spinlock lock;
  struct run *freelist;
} kmem;

struct {
    struct spinlock lock;
    int count[(PHYSTOP - KERNBASE) / PGSIZE]; // 物理页引用计数
} refcount;

// 辅助函数：获取物理页索引
uint64 pa2index(uint64 pa) {
    if(pa < KERNBASE || pa >= PHYSTOP)
        panic("pa2index: invalid pa");
    return (pa - KERNBASE) / PGSIZE;
}

void incref(void *pa) {
    if((uint64)pa % PGSIZE != 0 || (uint64)pa < KERNBASE || (uint64)pa >= PHYSTOP)
        panic("incref: invalid pa");
    
    acquire(&refcount.lock);
    refcount.count[pa2index((uint64)pa)]++;
    release(&refcount.lock);
}

void decref(void *pa) {
    if((uint64)pa % PGSIZE != 0 || (uint64)pa < KERNBASE || (uint64)pa >= PHYSTOP)
        panic("decref: invalid pa");
    
    acquire(&refcount.lock);
    refcount.count[pa2index((uint64)pa)]--;
    release(&refcount.lock);
}

void
kinit()
{
  initlock(&kmem.lock, "kmem");
  initlock(&refcount.lock, "refcount");
  
  // 初始化引用计数数组为0
  acquire(&refcount.lock);
  for(int i = 0; i < (PHYSTOP - KERNBASE) / PGSIZE; i++) {
      refcount.count[i] = 0;
  }
  release(&refcount.lock);
  
  freerange(end, (void*)PHYSTOP);
}

void
freerange(void *pa_start, void *pa_end)
{
  char *p;
  p = (char*)PGROUNDUP((uint64)pa_start);
  for(; p + PGSIZE <= (char*)pa_end; p += PGSIZE) {
    // 在初始化阶段，直接设置引用计数为1然后释放
    acquire(&refcount.lock);
    refcount.count[pa2index((uint64)p)] = 1;
    release(&refcount.lock);
    kfree(p);
  }
}

// Free the page of physical memory pointed at by pa,
// which normally should have been returned by a
// call to kalloc().  (The exception is when
// initializing the allocator; see kinit above.)
void
kfree(void *pa)
{
  struct run *r;

  if(((uint64)pa % PGSIZE) != 0 || (char*)pa < end || (uint64)pa >= PHYSTOP)
    panic("kfree");

  // 先检查引用计数
  acquire(&refcount.lock);
  int ref = refcount.count[pa2index((uint64)pa)];
  if (ref > 1) {
      // 还有引用，只减少计数不释放
      refcount.count[pa2index((uint64)pa)] = ref - 1;
      release(&refcount.lock);
      return;
  } else if (ref == 1) {
      // 最后一个引用，准备释放
      refcount.count[pa2index((uint64)pa)] = 0;
      release(&refcount.lock);
  } else {
      // ref == 0 或负数，错误
      release(&refcount.lock);
      panic("kfree: refcount error");
  }

  // Fill with junk to catch dangling refs.
  memset(pa, 1, PGSIZE);

  r = (struct run*)pa;

  acquire(&kmem.lock);
  r->next = kmem.freelist;
  kmem.freelist = r;
  release(&kmem.lock);
}

// Allocate one 4096-byte page of physical memory.
// Returns a pointer that the kernel can use.
// Returns 0 if the memory cannot be allocated.
void *
kalloc(void)
{
  struct run *r;

  acquire(&kmem.lock);
  r = kmem.freelist;
  if(r)
    kmem.freelist = r->next;
  release(&kmem.lock);

  if(r) {
    memset((char*)r, 5, PGSIZE); // fill with junk
    
    // 设置引用计数为1
    acquire(&refcount.lock);
    refcount.count[pa2index((uint64)r)] = 1;
    release(&refcount.lock);
  }
  
  return (void*)r;
}
